# 由参考值求当前信噪比
from typing import Dict, List
import json
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
from apps.cpm.rlm_context import RlmContext
from apps.cpm.cpm_engine import CpmEngine
from core.radar_core import RadarCore

class RefSnrAgent(object):
    def __init__(self):
        self.name = 'apps.cpm.agent.ref_snr_agent.RefSnrAgent'
    
    @staticmethod
    def run(context:RlmContext) -> Dict:
        
        '''
        作为整体任务第一步时，以用户自然语言为输入，输出为本步完成后的结果。
        '''
        observes = RefSnrAgent._observe(context=context)
        plans = RefSnrAgent._reason(observes=observes, context=context)
        return RefSnrAgent._act(observes=observes, context=context, plans=plans)

    @staticmethod
    def _observe(context:RlmContext) -> Dict:
        observes = {}
        prompt = RefSnrAgent._gen_params_parse_prompt(user_query=context.current_input)
        resp = CpmEngine.infer(query=prompt)
        print(f'resp json: \n{resp[8:-3]};')
        params = json.loads(resp[8:-3])
        context.states['params'].update(params)
        observes.update(params)
        return observes
    
    @staticmethod
    def _reason(observes:Dict, context: RlmContext) -> List:
        plans = []
        return plans
    
    @staticmethod
    def _act(observes:Dict, context: RlmContext, plans: List) -> Dict:
         # ---------------------------
        # 1. 参考参数设置
        # ---------------------------
        # 参考条件
        Rref = float(observes['Rref'])        # 参考距离 (m)
        tau_ref = float(observes['tau_ref'])  # 参考脉宽 (s)
        SNRref_db = float(observes['SNRref_db'])    # 参考SNR (dB)
        Sigmaref = float(observes['Sigmaref'])     # 参考RCS (m²)
        Lossp_db = float(observes['Lossp_db'])     # 处理损耗 (dB)
        # 当前条件
        R = float(observes['R'])        # 当前距离 (m)
        Sigma = float(observes['Sigma'])        # 当前RCS (m²)
        tau = float(observes['tau'])      # 当前脉宽 (s)
        snr_db = RadarCore.snr_by_ref(SNRref_db, Lossp_db, Rref, R, tau_ref, tau, Sigmaref, Sigma)
        # ---------------------------
        # 4. 结果验证
        # ---------------------------
        # 验证计算合理性
        assert not np.isinf(snr_db), "计算结果为无穷大，检查输入参数"
        assert not np.isnan(snr_db), "计算结果为NaN，检查输入参数"
        # ---------------------------
        # 5. 结果输出
        # ---------------------------
        print("\n=== 雷达参考距离计算结果 v0.0.1 ===")
        print(f"参考条件：")
        print(f"  参考距离 Rref = {Rref/1e3:.1f} km")
        print(f"  参考脉宽 tau_ref = {tau_ref*1e6:.1f} μs")
        print(f"  参考SNR SNRref = {SNRref_db:.1f} dB")
        print(f"  参考RCS Sigmaref = {Sigmaref} m²")
        print(f"  处理损耗 Lossp = {Lossp_db:.1f} dB\n")
        print(f"当前条件：")
        print(f"  当前距离 R = {R/1e3:.1f} km")
        print(f"  当前RCS Sigma = {Sigma} m²")
        print(f"  当前脉宽 tau = {tau*1e6:.1f} μs\n")
        print(f"计算结果：")
        # print(f"  线性域SNR = {snr_linear:.4e}")
        print(f"  dB域SNR = {snr_db:.2f} dB")
        # 规划执行步骤
        context.states['params'].update(observes)
        return {'result': 0, 'params': {}}
    
    @staticmethod
    def _gen_params_parse_prompt(user_query:str) -> str:
        template = '''
# 角色
假设你是资深雷达系统设计师

# 任务
你需要根据用户输入，提取出所有的参数，并转换为标准单位，以json格式输出。由于参数很多，而且参数表述和单位形式多样，你需要按照下面的步骤执行。

# 要求
你只需要给出json结果，不需要生成代码或进行解释。

# 步骤
1. 参考距离Rref：单位可能为m（米）、Km（千米），如果是Km需要将值乘以e3，提取值为v1；
2. 参考脉宽tau_ref：名称还可能为参考脉冲宽度，单位为s（秒）、ms（毫秒）、us（微秒），如果是ms时数值乘以e-3，如果是us时数值乘以e-6，提取值为v2；
3. 参考信噪比SNRref_db：名称还可能为参考SNR，如果数值没带dB，需要转为dB，公式为：dB=10*np.log10(SNRref_db)，提取值为v3；
4. 参考雷达截面积Sigmaref：名称还可能为参考RCS，单位为$m^{2}$，提取值为v4；
5. 处理损耗Lossp_db：名称还可能为损耗，如果数值没带dB，需要转为dB，公式为：dB=10*np.log10(Lossp_db)，提取值为v5；
6. 当前距离R：名称还可能为距离，单位可能为m（米）、Km（千米），如果是Km需要将值乘以e3，提取值为v6；
7. 雷达截面积Sigma：名称还可能为RCS、当前截面积、当前RCS，单位为$m^{2}$，提取值为v7；
8. 脉冲宽度tau：名称还可能为当前脉冲宽度、脉宽、当前脉宽，单位为s（秒）、ms（毫秒）、us（微秒），如果是ms时数值乘以e-3，如果是us时数值乘以e-6，提取值为v8；
9. 生成json结果：
{
    "Rref": v1,
    "tau_ref": v2,
    "SNRref_db": v3,
    "Sigmaref": v4,
    "Lossp_db": v5,
    "R": v6,
    "Sigma": v7,
    "tau": v8
}

# 示例
输入：某雷达参考距离为86Km，参考脉宽为0.1us，参考信噪比为20.0dB，参考雷达截面积为0.1$m^{2}$，处理损耗为2.0dB，当前距离为120Km，当前RCS为0.2$m^{2}$，当前脉冲宽度为0.1us，计算当前的信噪比。
输出：
{
    "Rref": 86e3,
    "tau_ref": 0.1e-6,
    "SNRref_db": 20.0,
    "Sigmaref": 0.1,
    "Lossp_db": 2.0,
    "R": 120e3,
    "Sigma": 0.2,
    "tau": 0.1e-6
}

# 当前输入
'''
        prompt = f'{template}{user_query}'
        return prompt